# Copyright (c) Microsoft Corporation.
# Licensed under the MIT license.
# Adapted from https://github.com/MARIO-Math-Reasoning/Super_MARIO
"""
MCTS（蒙特卡洛树搜索）算法实现

本模块实现了基于蒙特卡洛树搜索的推理算法，用于解决数学推理问题。
MCTS 通过选择、扩展、模拟和反向传播四个阶段来构建和优化搜索树。

主要功能:
- 实现 MCTS 算法的核心流程
- 支持节点选择和扩展策略
- 提供价值估计和反向传播机制
- 集成代码执行和错误处理
- 支持中间结果记录和评估

核心组件:
- MCTS: 主要的 MCTS 算法实现类
- 节点选择: 基于 PUCT 公式的子节点选择
- 节点扩展: 根据 LLM 输出创建新的子节点
- 价值更新: 递归更新节点价值和访问次数
"""
# 导入必要的库和模块
from __future__ import annotations  # 启用延迟注解，支持前向引用类型提示
from termcolor import colored  # 终端彩色输出库，用于调试信息的可视化显示
from typing import Dict, Any, Optional, Type, List, Tuple, Callable, Union  # 类型提示工具，提供静态类型检查支持
from pydantic import field_validator  # Pydantic 字段验证装饰器，用于配置参数验证
from vllm.outputs import CompletionOutput, RequestOutput  # VLLM 输出类型，处理语言模型的生成结果
from rstar_deepthink.agents.utils import math_equiv as is_equiv  # 数学等价性判断函数，用于评估答案正确性
from rstar_deepthink.nodes.base_node import BaseNode  # 搜索树节点基类
from rstar_deepthink.nodes import MCTSNode  # MCTS 专用节点类，包含 MCTS 算法所需的统计信息
from rstar_deepthink.constants import (  # 项目常量定义
    TOO_MANY_CODE_ERRORS,  # 连续代码错误过多的终止标记
    TOO_MANY_STEPS,       # 步数超限的终止标记
    NO_VALID_CHILD,       # 无有效子节点的终止标记
    CODE_END,             # 代码结束标记
)
from .tree import BaseTree, code_execution  # 搜索树基类和代码执行函数
from .beam_search import BS  # 束搜索算法基类


class MCTS(BS):
    """
    蒙特卡洛树搜索（MCTS）算法实现类
    
    继承自 BS（Beam Search）类，实现了完整的 MCTS 算法流程。
    MCTS 通过四个主要阶段来优化搜索过程：
    1. 选择（Selection）：从根节点开始，选择最有前途的路径
    2. 扩展（Expansion）：在选中的叶节点添加新的子节点
    3. 模拟（Simulation）：评估新节点的价值
    4. 反向传播（Backpropagation）：将评估结果向上传播
    
    属性:
        search_node: 当前搜索的节点
        intermediate_metric: 中间结果指标，用于记录搜索过程中的统计信息
    
    核心方法:
    - selection: 节点选择策略
    - expand_node: 节点扩展
    - eval_final_answer: 最终答案评估
    - select_next_step: 选择下一步操作
    """
    search_node: Type[BaseNode] = None

    # 中间结果指标，用于记录和分析搜索过程
    intermediate_metric: Dict = {
        "question": "",          # 问题文本
        "gt": "",               # 标准答案（ground truth）
        "answers": [],          # 候选答案列表
        "judgements": [],       # 答案正确性判断
        "value_estimate": [],   # 价值估计列表
        "rollout_indexs": [],   # 推演索引列表
    }

    @field_validator("config")
    def validate_config(cls, cfg: Any):
        """
        验证 MCTS 配置参数
        
        确保配置参数符合 MCTS 算法的要求，包括模式设置和停止条件。
        
        参数:
            cfg: 配置对象
            
        返回:
            验证后的配置对象
            
        异常:
            ValueError: 当配置参数不符合要求时抛出
        """
        BaseTree.validate_config(cfg)
        if not (cfg.mode == "mcts" or cfg.mode == "mcts_wo_code"):
            raise ValueError(f"Wrong value for config mode.")
        if cfg.stop is None:
            raise ValueError(f"Wrong value for config stop, cannot be None")
        return cfg

    def create_node(self, parent: Optional[Type[MCTSNode]] = None) -> Type[MCTSNode]:
        """
        创建新的 MCTS 节点
        
        根据父节点创建新的 MCTSNode 实例，设置必要的状态键和 PUCT 参数。
        
        参数:
            parent: 父节点，如果为 None 则创建根节点
            
        返回:
            Type[MCTSNode]: 新创建的 MCTS 节点
        """
        return MCTSNode(
            parent=parent, 
            additional_state_keys=self.NODE_KEYS,
            c_puct=self.config.c_puct,
        )


    def selection(self, from_root=False) -> Optional[Type[MCTSNode]]:
        """
        MCTS 选择阶段：选择最有前途的节点进行扩展
        
        从指定的起始节点开始，使用 PUCT 公式选择最优的子节点路径，
        直到找到一个可以扩展的叶节点。这是 MCTS 算法的第一个阶段。
        
        参数:
            from_root: 是否从根节点开始选择，默认从当前搜索节点开始
            
        返回:
            Optional[Type[MCTSNode]]: 选中的节点，如果没有可选节点则返回 None
            
        选择策略:
        1. 确定起始节点（根节点或当前搜索节点）
        2. 如果节点有子节点或已终止，尝试选择最优子节点
        3. 如果所有子节点都已终止，标记当前节点为终止状态
        """
        if from_root:
            start_node = self.root
        else:
            start_node = self.search_node
        
        # 选择子节点
        node = start_node
        if node is None: 
            return None
            
        if node.has_children() or node.is_terminal:
            # 鼓励探索，从非终止子节点中选择
            next_node = self.select_child(node)
            if next_node is None:  # 如果返回 None，说明所有子节点都已终止
                node.is_terminal = True
            node = next_node
            
        return None if (node is None or node.is_terminal) else node

    def select_child(self, node: Type[MCTSNode]) -> Optional[Type[MCTSNode]]:
        """
        使用 PUCT 公式选择最优子节点
        
        根据 PUCT（Polynomial Upper Confidence Trees）公式计算每个子节点的价值，
        选择价值最高的子节点。PUCT 平衡了利用（exploitation）和探索（exploration）。
        
        参数:
            node: 父节点
            
        返回:
            Optional[Type[MCTSNode]]: 选中的子节点，如果没有可选子节点则返回 None
            
        选择策略:
        1. 跳过所有已终止的子节点
        2. 计算每个子节点的 PUCT 值
        3. 选择 PUCT 值最高的子节点
        4. 如果有多个相同最高值，选择第一个
        """
        best_value = -float("inf")
        best_childs = []

        for child in node.children:
            if child.is_terminal:
                continue  # 跳过已终止的子节点
                
            puct_value = child.puct()
            if puct_value == best_value:
                best_childs.append(child)
            elif puct_value > best_value:
                best_value = puct_value
                best_childs = [child]

        # 注释：可以使用随机选择来增加多样性
        # return random.choice(best_childs) if best_childs else None
        return best_childs[0] if best_childs else None

    def expand_node(self, outputs: List[CompletionOutput], node: Type[MCTSNode]) -> None:
        """
        MCTS 扩展阶段：为选中的节点创建子节点
        
        根据 LLM 的输出结果，为指定节点创建多个子节点。
        这是 MCTS 算法的第二个阶段，用于扩展搜索树。
        
        参数:
            outputs: LLM 生成的输出列表
            node: 要扩展的父节点
            
        工作流程:
        1. 遍历所有 LLM 输出
        2. 解析每个输出的内容和停止原因
        3. 为每个有效输出创建对应的子节点
        """
        for idx, output in enumerate(outputs):
            # 确保停止原因不为空
            if not output.stop_reason: 
                output.stop_reason = ""
            # 解析输出内容
            step_result, parser_result = self.step_unwrap(output.text + output.stop_reason)
            # 创建子节点
            self.create_child(step_result, parser_result, node, idx)

    def create_child(
        self, 
        step_result: str, 
        parser_result: Dict[str, str], 
        node: Type[MCTSNode],
        idx: int,
    ) -> None:
        """
        为指定父节点创建子节点
        
        根据 LLM 生成的步骤结果和解析结果，创建新的子节点并设置其状态。
        处理不同类型的解析结果：最终答案、代码执行、普通文本等。
        
        参数:
            step_result: LLM 生成的步骤文本
            parser_result: 解析后的结构化结果，包含 action、action_input、final_answer 等字段
            node: 父节点
            idx: 子节点索引（当前未使用）
            
        处理逻辑:
        1. 如果解析失败 (parser_result is None)：标记为终端节点，设置 NO_VALID_CHILD
        2. 如果包含最终答案：标记为终端节点，评估答案正确性
        3. 如果包含代码动作：执行代码，处理观察结果和错误
        4. 检查深度限制，超过最大深度则终止
        """
        new_node = self.create_node(parent=node)
        parent_child_count = len(node.children)
        new_node.tag = f"{node.tag}.{parent_child_count + 1}"
        new_node.depth = node.depth + 1

        if parser_result is None:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = NO_VALID_CHILD
            self.eval_final_answer(new_node)
        elif parser_result["final_answer"]:
            new_node.is_terminal = True
            new_node.state["text"] = step_result
            new_node.state["final_answer"] = parser_result["final_answer"]
            self.eval_final_answer(new_node)
        elif parser_result["action"]:
            observation = code_execution(node, parser_result)
            new_node.state["action"] = parser_result["action"]
            new_node.state["action_input"] = parser_result["action_input"]
            new_node.state["observation"] = observation
            if CODE_END in parser_result["action_input"]:
                observation = self.obs_wrap(observation)
                new_node.state["text"] = f"{step_result}{self.config.step_delim}{observation}"
            else:
                new_node.state["text"] = step_result
                
            if "error" in observation.lower():
                new_node.consecutive_errors = node.consecutive_errors + 1
                if new_node.consecutive_errors >= self.config.errors_threshold:
                    observation = self.obs_wrap(observation)
                    step_result = step_result + CODE_END if CODE_END not in step_result else step_result
                    new_node.state["text"] = f"{step_result}{self.config.step_delim}{observation}"
                    new_node.is_terminal = True
                    new_node.state["final_answer"] = TOO_MANY_CODE_ERRORS
                    self.eval_final_answer(new_node)
        else:
            new_node.state["text"] = step_result

        if not new_node.is_terminal and new_node.depth > self.config.max_depth:
            new_node.is_terminal = True
            new_node.state["final_answer"] = TOO_MANY_STEPS
            self.eval_final_answer(new_node)

        node.children.append(new_node)

    def eval_final_answer(self, node: Type[MCTSNode]) -> None:
        if node.state["final_answer"] in [NO_VALID_CHILD, TOO_MANY_STEPS, TOO_MANY_CODE_ERRORS]:
            # if the final answer is not valid, update the node with negative reward
            node.update(self.config.negative_reward)
            return 
        
        if self.config.is_sampling:
            final_answer = node.state["final_answer"]
            correct = is_equiv(self.ground_truth, final_answer)
            node.update_recursive(self.config.positive_reward if correct else self.config.negative_reward, self.root)
        else:
            # just append the node to candidate_nodes, will update the value in select_next_step()
            self.candidate_nodes.append(node)


    def record_intermediate_metric(self, answer: str, value_estimate: float) -> None:
        """
        记录中间评估指标
        
        在 MCTS 搜索过程中记录每个 rollout 的最佳答案和对应的价值估计。
        每个 rollout 只保留价值估计最高的答案，用于后续分析和调试。
        
        参数:
            answer: 候选答案
            value_estimate: 该答案的价值估计
            
        功能:
        1. 记录问题和标准答案
        2. 为每个 rollout 保留最高价值的答案
        3. 计算答案的正确性判断
        4. 维护 rollout 索引和对应的评估结果
        """
        self.intermediate_metric["question"] = self.question
        self.intermediate_metric["gt"] = self.ground_truth
        # each rollout retains the answer with the highest value_estimate
        # Check if the rollout's answer is already in the list
        if self.intermediate_metric["rollout_indexs"] and self.rollout_idx in self.intermediate_metric["rollout_indexs"]:
            # Find the index of the existing rollout
            index = self.intermediate_metric["rollout_indexs"].index(self.rollout_idx)
            if value_estimate > self.intermediate_metric["value_estimate"][index]:
                self.intermediate_metric["answers"][index] = answer
                self.intermediate_metric["judgements"][index] = is_equiv(self.ground_truth, answer)
                self.intermediate_metric["value_estimate"][index] = value_estimate
        else:
            # If the rollout's answer is not in the list, add it
            self.intermediate_metric["answers"].append(answer)
            self.intermediate_metric["judgements"].append(is_equiv(self.ground_truth, answer))
            self.intermediate_metric["value_estimate"].append(value_estimate)
            self.intermediate_metric["rollout_indexs"].append(self.rollout_idx)


    def select_next_step(self, outputs=None, from_root: bool = False) -> None:
        """
        选择下一步要扩展的节点
        
        MCTS 算法的选择阶段，处理候选节点的价值更新和下一个搜索节点的选择。
        根据价值模型的输出更新节点价值，并使用 UCB 公式选择最优节点。
        
        参数:
            outputs: 价值模型的输出结果列表
            from_root: 是否从根节点开始选择
            
        工作流程:
        1. 处理候选节点的价值估计
        2. 对终端节点进行递归价值更新
        3. 对中间节点进行局部价值更新
        4. 收集有效的最终答案节点
        5. 使用 selection 方法选择下一个扩展节点
        """
        self.search_node = self.current_nodes[0] if self.current_nodes else None
        self.current_nodes = []
        if outputs:
            for candidate_node, output in zip(self.candidate_nodes, outputs):
                if candidate_node.is_terminal and self.config.is_sampling:
                    continue
                value_estimate = output.value_estimate if output.value_estimate is not None else self.config.negative_reward
                if output.value_estimate is None:
                    candidate_node.is_terminal = True

                # backup
                if candidate_node.is_terminal and candidate_node.state["final_answer"]:
                    # for terminal node: update_recursive
                    if candidate_node.state["final_answer"] in [NO_VALID_CHILD, TOO_MANY_STEPS, TOO_MANY_CODE_ERRORS]:
                        candidate_node.update(self.config.negative_reward)
                    else:
                        # save intermediate metric
                        self.record_intermediate_metric(answer=candidate_node.state["final_answer"], value_estimate=value_estimate)

                        candidate_node.update_recursive(value_estimate, self.root)
                else:
                    # for intermediate node: just update the value
                    if self.config.terminal_sample:
                        pass
                    else:
                        candidate_node.update(value_estimate)

                if self.__class__.is_valid_final_answer_node(candidate_node):
                    self.final_answer_nodes.append(candidate_node)
        selection_node = self.selection(from_root=from_root)
        if selection_node is not None:
            self.current_nodes.append(selection_node)
    
    
    def generate_next_step(self, outputs: List[RequestOutput]) -> None:
        """
        生成下一步的候选节点
        
        MCTS 算法的扩展阶段，根据 LLM 的输出为当前节点创建子节点。
        处理价值估计并准备候选节点列表供后续选择。
        
        参数:
            outputs: LLM 生成的输出结果列表
            
        工作流程:
        1. 清空候选节点列表
        2. 遍历当前节点和对应的输出
        3. 检查价值估计的有效性
        4. 调用 expand_node 创建子节点
        5. 如果启用叶节点价值更新，将新创建的子节点加入候选列表
        """
        self.candidate_nodes = []
        for current_node, output in zip(self.current_nodes, outputs):
            value_estimate = output.value_estimate
            if value_estimate is not None:  
                self.expand_node(output.outputs, current_node)
            else:
                value_estimate = self.config.negative_reward
                current_node.is_terminal = True

            if self.config.update_leaf_value:
                # if need update leaf node value, just append the node to candidate_nodes, will update the value in select_next_step()
                for value_node in current_node.children:
                    if value_node not in self.candidate_nodes and value_node.visit_count() < 1:
                        self.candidate_nodes.append(value_node) 
                    
                    
    def return_states(self) -> Dict[str, Union[Any, Dict[str, str]]]:
        """
        返回搜索树的完整状态信息
        
        遍历整个搜索树，收集所有节点的状态、价值、访问次数等信息。
        用于调试、可视化和结果分析。
        
        返回:
            Dict: 包含所有节点状态的字典，键为节点标签，值为节点的详细信息
                 包括：state（节点状态）、value（节点价值）、q_value（Q值）、visit_count（访问次数）
                 
        实现方式:
        使用广度优先搜索遍历整个树结构，确保所有节点都被包含在结果中
        """
        candidates = [self.root]
        states = {}
        while candidates:
            node = candidates.pop(0)
            states[node.tag] = node.state
            states[node.tag]["value"] = node.value
            states[node.tag]["q_value"] = node.q_value()
            states[node.tag]["visit_count"] = node.visit_count()
            if node.has_children():
                candidates.extend(node.children)
        return states
